import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk
import cv2
import numpy as np
import imgaug.augmenters as iaa

"""
数据增强可视化工具
对图片进行可视化增强效果，使研究人员对于数据增强参数的数值有一个感性认知
"""

############需要修改###################
# 指定图片和文本文件夹的路径
path = 'img/1.jpg'
#####################################

class ImageEnhancementApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Image Enhancement")
        
        # 初始化亮度和模糊调整值
        self.brightness_value = 1.0
        self.blur_value = 3
        
        # 创建滑块和标签
        self.slider_label1 = tk.Label(root, text="Brightness Adjustment")
        self.slider_label1.pack()
        # 滑块的起始值与最大值，步长 ####################
        self.brightness_slider = tk.Scale(root, from_=0.1, to=3.0, resolution=0.1, orient="horizontal", command=self.update_brightness)
        self.brightness_slider.set(1.0)
        self.brightness_slider.pack()
        
        self.slider_label2 = tk.Label(root, text="Blur Adjustment")
        self.slider_label2.pack()
        # 滑块的起始值与最大值，步长 ####################
        self.blur_slider = tk.Scale(root, from_=3, to=20, orient="horizontal", command=self.update_blur)
        self.blur_slider.set(1)
        self.blur_slider.pack()
        
        # 创建画布
        self.canvas = tk.Canvas(root)
        self.canvas.pack()
        
        # 加载指定图片
        self.image_path = path  # 指定图片路径
        self.image = cv2.imread(self.image_path)
        self.display_image()

    def display_image(self):
        # 将图像显示在画布上
        if self.image is not None:
            # 调整亮度
            image_enhanced = np.clip(self.image * self.brightness_value, 0, 255).astype(np.uint8)
            # 应用动态模糊
            image_enhanced = self.apply_dynamic_blur(image_enhanced, self.blur_value)
            # 将图像转换为RGB格式
            image_rgb = cv2.cvtColor(image_enhanced, cv2.COLOR_BGR2RGB)
            # 将图像转换为PIL格式
            image_pil = Image.fromarray(image_rgb)
            # 根据图像大小调整画布大小
            self.canvas.config(width=image_pil.width, height=image_pil.height)
            # 在画布上显示图像
            self.photo = ImageTk.PhotoImage(image_pil)
            self.canvas.create_image(0, 0, anchor=tk.NW, image=self.photo)

    def update_brightness(self, value):
        # 更新亮度调整值并重新显示图像
        self.brightness_value = float(value)
        self.display_image()
        
    def update_blur(self, value):
        # 更新模糊调整值并重新显示图像
        self.blur_value = int(value)
        self.display_image()

    def apply_dynamic_blur(self, image, blur_value):
        # 创建动态模糊增强器
        seq = iaa.Sequential([
            iaa.MotionBlur(k=(blur_value ), angle=(88,92))
        ])
        # 将图像转换为float32格式并调整亮度
        image = image.astype(np.float32) / 255.0
        # 应用动态模糊增强
        image_enhanced = seq(images=[image])[0]
        # 将图像转换回uint8格式
        return (image_enhanced * 255).astype(np.uint8)

if __name__ == "__main__":
    root = tk.Tk()
    app = ImageEnhancementApp(root)
    root.mainloop()